from llama_index.utils.workflow import draw_all_possible_flows,draw_all_possible_flows_mermaid

import asyncio
import random
 
from llama_index.core.workflow import (
    StartEvent,
    StopEvent,
    Workflow,
    step,
    Event,
)
from llama_index.utils.workflow import draw_all_possible_flows
 
class LoopEvent(Event):
    loop_output: str
 
class FirstEvent(Event):
    first_output: str
 
class SecondEvent(Event):
    second_output: str
 
 
class MyWorkflow(Workflow):
 
    @step
    async def step_one(self, ev: StartEvent | LoopEvent)-> FirstEvent | LoopEvent:
        """第一步"""
        if random.randint(0, 1) == 0:
            print("发生了不好的事情")
            return LoopEvent(loop_output="返回到第一步...")
        else:
            print("开始第一步...")
            return FirstEvent(first_output="完成第一步...")
 
    @step
    async def step_two(self, ev: FirstEvent)-> SecondEvent:
        """第二步"""
        print(ev.first_output)
        return SecondEvent(second_output="完成第二步...")
 
    @step
    async def step_three(self, ev: SecondEvent)-> StopEvent:
        """第三步"""
        print(ev.second_output)
        return StopEvent(result="工作流完成!")
 
# 运行工作流
async def run_workflow():
    w = MyWorkflow(timeout=10, verbose=False)
    result = await w.run(first_input="开始工作流...")
    print(result)
 
if __name__ == '__main__':
    asyncio.run(run_workflow())
 
# 工作流可视化工具
draw_all_possible_flows(MyWorkflow, filename="multi_step_workflow.html")

draw_all_possible_flows_mermaid(MyWorkflow, filename="multi_step_workflow.flowchart")